package com.abyss.state;


import org.apache.flink.api.common.restartstrategy.RestartStrategies;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.time.Time;
import org.apache.flink.api.common.typeinfo.TypeHint;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.runtime.state.filesystem.FsStateBackend;
import org.apache.flink.streaming.api.CheckpointingMode;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.CheckpointConfig;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.source.SourceFunction;

/**
 * 做一个自定义的source, source在keyby之前，那么用的状态就是OperatorState
 * source会累加生成数字
 */
public class OperatorStateDemo {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();



        env.enableCheckpointing(5000L, CheckpointingMode.EXACTLY_ONCE);
        env.setStateBackend(new FsStateBackend("file:///Users/abyss/Dev/toys/flink/H-flink-learn/src/main/resources/OperatorState"));
        env.setRestartStrategy(RestartStrategies.failureRateRestart(7, Time.minutes(1), Time.seconds(5)));
        env.getCheckpointConfig().enableExternalizedCheckpoints(CheckpointConfig.ExternalizedCheckpointCleanup.RETAIN_ON_CANCELLATION);
        DataStreamSource<Integer> source = env.addSource(new MySourceForOperatorState());

        source.print();
        env.execute();
    }

    /**
     * 自定义Source 累加生成数字 从0开始 每一次累加1 每一秒生成一个
     */
    public static class MySourceForOperatorState implements CheckpointedFunction, SourceFunction<Integer> {
        private ListStateDescriptor<Integer> listStateDescriptor =
                new ListStateDescriptor<>("operatorState", TypeInformation.of(new TypeHint<Integer>() {}));
        private ListState<Integer> listState;
        private int lastAccumulator = 0;
        private boolean isRun = true;
        @Override
        public void snapshotState(FunctionSnapshotContext context) throws Exception {
            this.listState.clear(); //清除旧数据
            this.listState.add(this.lastAccumulator);
        }

        @Override
        public void initializeState(FunctionInitializationContext context) throws Exception {
            listState = context.getOperatorStateStore().getListState(listStateDescriptor);

            for (int i: listState.get()) {
                this.lastAccumulator = i;
            }
        }

        @Override
        public void run(SourceContext<Integer> ctx) throws Exception {
            while (isRun) {
                this.lastAccumulator++;
                ctx.collect(this.lastAccumulator);
                Thread.sleep(100L);
                if (this.lastAccumulator % 10 == 0) {
                    System.out.println(1 / 0);
                }
            }
        }

        @Override
        public void cancel() {
            this.isRun = false;
        }
    }
}